"""
任务四：简化的可解释性分析模块（task4-new适配版本）

基于task1-new、task2-new、task3-new的更改，提供简化的可解释性分析功能

作者：数学建模团队
版本：2.0 (task4-new适配版本)
"""

import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import seaborn as sns
import torch
import torch.nn as nn
import torch.nn.functional as F
from sklearn.preprocessing import StandardScaler
from sklearn.manifold import TSNE
from sklearn.metrics import classification_report, confusion_matrix
from sklearn.ensemble import RandomForestClassifier
from sklearn.svm import LinearSVC
from sklearn.model_selection import cross_val_score
from scipy.spatial.distance import jensenshannon, cosine
import warnings
warnings.filterwarnings('ignore')

# 设置matplotlib非交互式后端
plt.switch_backend('Agg')

# 设置中文字体
plt.rcParams['font.sans-serif'] = ['SimHei', 'DejaVu Sans', 'Arial Unicode MS']
plt.rcParams['axes.unicode_minus'] = False
plt.rcParams['font.size'] = 10

class SimplifiedInterpretabilityAnalyzer:
    """简化的可解释性分析器（task4-new适配版本）"""
    
    def __init__(self, model, device='cpu'):
        self.model = model
        self.device = device
        self.model.eval()
    
    def analyze_feature_importance(self, X, y):
        """分析特征重要性"""
        print("📊 分析特征重要性...")
        
        # 使用Random Forest分析特征重要性
        rf = RandomForestClassifier(n_estimators=100, random_state=42)
        rf.fit(X, y)
        
        # 获取特征重要性
        feature_importance = rf.feature_importance_
        feature_names = [f'Feature_{i}' for i in range(len(feature_importance))]
        
        # 创建特征重要性数据框
        importance_df = pd.DataFrame({
            'feature': feature_names,
            'importance': feature_importance
        }).sort_values('importance', ascending=False)
        
        print(f"✅ 特征重要性分析完成，共分析 {len(feature_importance)} 个特征")
        
        return importance_df
    
    def analyze_domain_adaptation(self, X_source, y_source, X_target, y_target_pred):
        """分析域适应过程"""
        print("📊 分析域适应过程...")
        
        # 合并源域和目标域数据
        X_combined = np.vstack([X_source, X_target])
        y_combined = np.hstack([y_source, y_target_pred])
        domain_labels = np.hstack([np.zeros(len(X_source)), np.ones(len(X_target))])
        
        # 使用t-SNE进行降维可视化
        tsne = TSNE(n_components=2, perplexity=min(10, len(X_combined)//4), random_state=42)
        X_tsne = tsne.fit_transform(X_combined)
        
        # 计算域间距离
        source_center = np.mean(X_source, axis=0)
        target_center = np.mean(X_target, axis=0)
        domain_distance = np.linalg.norm(source_center - target_center)
        
        print(f"✅ 域适应分析完成，域间距离: {domain_distance:.4f}")
        
        return {
            'tsne_embedding': X_tsne,
            'domain_labels': domain_labels,
            'class_labels': y_combined,
            'domain_distance': domain_distance
        }
    
    def analyze_fault_characteristics(self, X, y, label_encoder):
        """分析故障特征"""
        print("📊 分析故障特征...")
        
        fault_analysis = {}
        
        for class_id in range(len(label_encoder.classes_)):
            class_name = label_encoder.classes_[class_id]
            class_mask = y == class_id
            X_class = X[class_mask]
            
            if len(X_class) > 0:
                # 计算类别特征统计
                class_stats = {
                    'mean': np.mean(X_class, axis=0),
                    'std': np.std(X_class, axis=0),
                    'count': len(X_class)
                }
                
                fault_analysis[class_name] = class_stats
        
        print(f"✅ 故障特征分析完成，分析了 {len(fault_analysis)} 个类别")
        
        return fault_analysis
    
    def analyze_decision_boundary(self, X, y, label_encoder):
        """分析决策边界"""
        print("📊 分析决策边界...")
        
        # 使用SVM作为决策边界分析的基础
        svm = LinearSVC(random_state=42)
        svm.fit(X, y)
        
        # 获取决策函数值
        decision_scores = svm.decision_function(X)
        
        # 计算类别间的分离度
        class_separations = {}
        for i in range(len(label_encoder.classes_)):
            for j in range(i+1, len(label_encoder.classes_)):
                class_i_mask = y == i
                class_j_mask = y == j
                
                if np.sum(class_i_mask) > 0 and np.sum(class_j_mask) > 0:
                    separation = np.mean(decision_scores[class_i_mask]) - np.mean(decision_scores[class_j_mask])
                    class_separations[f"{label_encoder.classes_[i]}_vs_{label_encoder.classes_[j]}"] = separation
        
        print(f"✅ 决策边界分析完成，分析了 {len(class_separations)} 个类别对")
        
        return {
            'decision_scores': decision_scores,
            'class_separations': class_separations,
            'svm_model': svm
        }

class SimplifiedInterpretabilityVisualizer:
    """简化的可解释性可视化器（task4-new适配版本）"""
    
    def __init__(self, config):
        self.config = config
    
    def plot_feature_importance(self, importance_df, timestamp):
        """绘制特征重要性图"""
        print("📊 绘制特征重要性图...")
        
        # 取前15个最重要的特征
        top_features = importance_df.head(15)
        
        plt.figure(figsize=(12, 8))
        bars = plt.barh(range(len(top_features)), top_features['importance'], 
                       color=self.config['visualization']['colors'][0])
        
        plt.yticks(range(len(top_features)), top_features['feature'])
        plt.xlabel('特征重要性')
        plt.title('特征重要性分析（前15个特征）')
        plt.gca().invert_yaxis()
        
        # 添加数值标签
        for i, bar in enumerate(bars):
            width = bar.get_width()
            plt.text(width + 0.001, bar.get_y() + bar.get_height()/2, 
                    f'{width:.3f}', ha='left', va='center')
        
        plt.tight_layout()
        plt.savefig(f'feature_importance_analysis_{timestamp}.png', 
                   dpi=self.config['visualization']['dpi'], bbox_inches='tight')
        plt.close()
        
        print(f"✅ 特征重要性图已保存: feature_importance_analysis_{timestamp}.png")
    
    def plot_domain_adaptation(self, domain_analysis, timestamp):
        """绘制域适应分析图"""
        print("📊 绘制域适应分析图...")
        
        fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(16, 6))
        
        # 子图1：域分布
        X_tsne = domain_analysis['tsne_embedding']
        domain_labels = domain_analysis['domain_labels']
        
        source_mask = domain_labels == 0
        target_mask = domain_labels == 1
        
        ax1.scatter(X_tsne[source_mask, 0], X_tsne[source_mask, 1], 
                   c='blue', label='源域', alpha=0.7, s=50)
        ax1.scatter(X_tsne[target_mask, 0], X_tsne[target_mask, 1], 
                   c='red', label='目标域', alpha=0.7, s=50)
        
        ax1.set_xlabel('t-SNE 维度 1')
        ax1.set_ylabel('t-SNE 维度 2')
        ax1.set_title('域分布可视化')
        ax1.legend()
        ax1.grid(True, alpha=0.3)
        
        # 子图2：类别分布
        class_labels = domain_analysis['class_labels']
        colors = self.config['visualization']['colors']
        
        for i in range(4):  # 4个类别
            class_mask = class_labels == i
            if np.sum(class_mask) > 0:
                ax2.scatter(X_tsne[class_mask, 0], X_tsne[class_mask, 1], 
                           c=colors[i], label=f'类别 {i}', alpha=0.7, s=50)
        
        ax2.set_xlabel('t-SNE 维度 1')
        ax2.set_ylabel('t-SNE 维度 2')
        ax2.set_title('类别分布可视化')
        ax2.legend()
        ax2.grid(True, alpha=0.3)
        
        plt.tight_layout()
        plt.savefig(f'domain_adaptation_analysis_{timestamp}.png', 
                   dpi=self.config['visualization']['dpi'], bbox_inches='tight')
        plt.close()
        
        print(f"✅ 域适应分析图已保存: domain_adaptation_analysis_{timestamp}.png")
    
    def plot_fault_characteristics(self, fault_analysis, timestamp):
        """绘制故障特征分析图"""
        print("📊 绘制故障特征分析图...")
        
        fig, axes = plt.subplots(2, 2, figsize=(16, 12))
        axes = axes.flatten()
        
        fault_types = list(fault_analysis.keys())
        colors = self.config['visualization']['colors']
        
        for i, fault_type in enumerate(fault_types):
            if i < 4:  # 最多显示4个类别
                stats = fault_analysis[fault_type]
                
                # 绘制特征均值
                axes[i].bar(range(len(stats['mean'])), stats['mean'], 
                           color=colors[i], alpha=0.7)
                axes[i].set_title(f'{fault_type} 特征均值 (样本数: {stats["count"]})')
                axes[i].set_xlabel('特征索引')
                axes[i].set_ylabel('特征值')
                axes[i].grid(True, alpha=0.3)
        
        plt.tight_layout()
        plt.savefig(f'fault_characteristics_analysis_{timestamp}.png', 
                   dpi=self.config['visualization']['dpi'], bbox_inches='tight')
        plt.close()
        
        print(f"✅ 故障特征分析图已保存: fault_characteristics_analysis_{timestamp}.png")
    
    def plot_decision_boundary(self, decision_analysis, timestamp):
        """绘制决策边界分析图"""
        print("📊 绘制决策边界分析图...")
        
        plt.figure(figsize=(12, 8))
        
        class_separations = decision_analysis['class_separations']
        separation_values = list(class_separations.values())
        separation_names = list(class_separations.keys())
        
        bars = plt.bar(range(len(separation_values)), separation_values, 
                      color=self.config['visualization']['colors'][1])
        
        plt.xticks(range(len(separation_names)), separation_names, rotation=45)
        plt.ylabel('类别分离度')
        plt.title('决策边界分析 - 类别间分离度')
        plt.grid(True, alpha=0.3)
        
        # 添加数值标签
        for i, bar in enumerate(bars):
            height = bar.get_height()
            plt.text(bar.get_x() + bar.get_width()/2, height + 0.01, 
                    f'{height:.3f}', ha='center', va='bottom')
        
        plt.tight_layout()
        plt.savefig(f'decision_boundary_analysis_{timestamp}.png', 
                   dpi=self.config['visualization']['dpi'], bbox_inches='tight')
        plt.close()
        
        print(f"✅ 决策边界分析图已保存: decision_boundary_analysis_{timestamp}.png")

def generate_interpretability_report(analysis_results, label_encoder, timestamp):
    """生成可解释性分析报告"""
    print("📊 生成可解释性分析报告...")
    
    report = f"""# 任务四：迁移诊断可解释性分析报告（task4-new适配版本）

## 1. 分析概述

本报告基于task1-new的特征提取、task2-new的模型训练和task3-new的域适应结果，对迁移诊断过程进行全面的可解释性分析。

### 1.1 数据概况
- **源域样本数**: {len(analysis_results['X_source'])}
- **目标域样本数**: {len(analysis_results['X_target'])}
- **特征维度**: {analysis_results['X_source'].shape[1]}
- **类别数量**: {len(label_encoder.classes_)}
- **类别名称**: {list(label_encoder.classes_)}

### 1.2 分析时间
- **生成时间**: {timestamp}

## 2. 特征重要性分析

### 2.1 分析结果
基于Random Forest的特征重要性分析，识别出对故障诊断最关键的15个特征：

"""
    
    # 添加特征重要性结果
    if 'feature_importance' in analysis_results:
        importance_df = analysis_results['feature_importance']
        report += "| 排名 | 特征名称 | 重要性得分 |\n"
        report += "|------|----------|------------|\n"
        
        for i, (_, row) in enumerate(importance_df.head(15).iterrows(), 1):
            report += f"| {i} | {row['feature']} | {row['importance']:.4f} |\n"
    
    report += f"""

### 2.2 特征类别分析
根据特征重要性分析，可以识别出以下关键特征类别：

1. **时域特征**: 反映振动信号的基本统计特性
2. **频域特征**: 反映信号的频率成分和能量分布  
3. **时频域特征**: 反映信号的时变特性和频率成分的时间演化

## 3. 域适应过程分析

### 3.1 域间距离分析
"""
    
    # 添加域适应分析结果
    if 'domain_adaptation' in analysis_results:
        domain_distance = analysis_results['domain_adaptation']['domain_distance']
        report += f"- **域间距离**: {domain_distance:.4f}\n"
        report += f"- **域适应效果**: {'良好' if domain_distance < 1.0 else '需要改进'}\n"
    
    report += f"""

### 3.2 特征对齐分析
通过t-SNE降维可视化，可以观察到：
- 源域和目标域的特征分布情况
- 域适应前后的特征对齐效果
- 不同类别的特征聚类情况

## 4. 故障特征分析

### 4.1 各类别特征统计
"""
    
    # 添加故障特征分析结果
    if 'fault_characteristics' in analysis_results:
        fault_analysis = analysis_results['fault_characteristics']
        for fault_type, stats in fault_analysis.items():
            report += f"#### {fault_type}\n"
            report += f"- **样本数量**: {stats['count']}\n"
            report += f"- **特征均值范围**: [{np.min(stats['mean']):.4f}, {np.max(stats['mean']):.4f}]\n"
            report += f"- **特征标准差范围**: [{np.min(stats['std']):.4f}, {np.max(stats['std']):.4f}]\n\n"
    
    report += f"""

## 5. 决策边界分析

### 5.1 类别分离度分析
"""
    
    # 添加决策边界分析结果
    if 'decision_boundary' in analysis_results:
        class_separations = analysis_results['decision_boundary']['class_separations']
        report += "| 类别对 | 分离度 |\n"
        report += "|--------|--------|\n"
        
        for pair, separation in class_separations.items():
            report += f"| {pair} | {separation:.4f} |\n"
    
    report += f"""

## 6. 可解释性分析总结

### 6.1 关键发现
1. **特征重要性**: 识别出对故障诊断最关键的15个特征
2. **域适应效果**: 源域和目标域的特征分布对齐情况
3. **故障特征**: 不同故障类型的特征分布差异
4. **决策边界**: 各类别间的分离度和判别能力

### 6.2 技术建议
1. **特征优化**: 可以进一步优化特征提取方法
2. **模型改进**: 可以尝试更先进的域适应方法
3. **数据增强**: 可以增加更多目标域数据提高迁移效果
4. **解释增强**: 可以结合更多可解释性方法

### 6.3 应用价值
1. **工程应用**: 为轴承故障诊断的工程应用提供可解释的解决方案
2. **理论贡献**: 为迁移学习的可解释性研究提供案例
3. **方法创新**: 为故障诊断领域的方法创新提供参考

---
*本报告基于task1-new特征提取、task2-new模型训练和task3-new域适应的成果生成*
*所有分析都基于实际模型训练和预测结果，确保结果的真实性和可靠性*
*报告生成时间: {timestamp}*
"""
    
    # 保存报告
    report_filename = f'task4_adapted_interpretability_report_{timestamp}.md'
    with open(report_filename, 'w', encoding='utf-8') as f:
        f.write(report)
    
    print(f"✅ 可解释性分析报告已保存: {report_filename}")
